from fastapi import Depends
import typing as T
import numpy as np
from pymilvus import MilvusClient
from ...init.Global import EmbeddingSession


class _RAGMapper:
    session: MilvusClient = None

    def __call__(self, session: EmbeddingSession):
        self.session = session
        return self

    async def searchByQueryAndKBId(
        self,
        query_embeddiing: list[float] | np.ndarray,
        knowledgeBaseId: int,
        top_n: int = 1,
        collection_name: str = "default",
        exclude_image: bool = True,
    ):
        result = (
            self.session.search(
                collection_name,
                [query_embeddiing],
                filter=(
                    f"knowledgeBaseId == {knowledgeBaseId}"
                    if not exclude_image
                    else f"knowledgeBaseId == {knowledgeBaseId} and metadata['type'] != 'image'"
                ),
                output_fields=["content", "knowledgeBaseId", "metadata", "userId"],
                limit=top_n,
                anns_field="embedding",
                search_params={
                    "params": {
                        "nprobe": 16,
                    }
                },
            )
        )[0]
        return result


RAGMapper = T.Annotated[_RAGMapper, Depends(_RAGMapper())]
